import re
import pytesseract
from pdf2image import convert_from_path
from collections import defaultdict
from PIL import ImageEnhance, ImageFilter

from Contexts import pdf_to_text

# 配置OCR路径（根据实际环境修改）
pytesseract.pytesseract.tesseract_cmd = r'E:\TesseractOcr\tesseract.exe'  # Tesseract路径
poppler_path = r'E:\poppler\poppler\poppler-24.08.0\Library\bin'  # Windows需配置，Mac/Linux删除

domain_keywords = {
    "直肠癌", "淋巴结转移", "智能诊断", "CT影像", "肿瘤分割",
    "特征提取", "分类模型", "Dice系数", "随机森林", "Mask RCNN"
}


# --------------------------
# OCR预处理函数
# --------------------------
# def preprocess_image(image):
#     """图像增强处理"""
#     # 转为灰度图
#     image = image.convert('L')
#     # 增强对比度
#     image = ImageEnhance.Contrast(image).enhance(2.0)
#     # 锐化处理
#     image = image.filter(ImageFilter.SHARPEN)
#     # 二值化降噪
#     image = image.point(lambda x: 0 if x < 180 else 255, '1')
#     return image
#
#
# def ocr_pdf_to_text(pdf_path):
#     """将扫描版PDF转换为结构化文本"""
#     pages_text = []
#     images = convert_from_path(pdf_path, dpi=300, poppler_path=poppler_path)
#
#     for img in images:
#         # 图像预处理
#         processed_img = preprocess_image(img)
#         # OCR识别（中英混合模式）
#         text = pytesseract.image_to_string(processed_img, lang='chi_sim+eng')
#         pages_text.append(text)
#
#     return pages_text


# --------------------------
# 核心标题提取逻辑（优化版）
# --------------------------
def extract_title_from_ocr_text(pages_text):
    """从OCR文本中提取标题"""
    # 合并前3页文本（根据需求调整）
    merged_text = "\n".join(pages_text[:3])

    # 优化后的标题正则（兼容OCR可能的换行错误）
    title_pattern = re.compile(
        r"(?<![\d.])"  # 排除编号开头
        r"(?:[《]?[^\n\r\s]{2,}[\u4e00-\u9fa5][^\n\r]{8,}?)"  # 核心标题段
        r"(?:\s*[：\-—]\s*[^\n\r\s]{2,}[\u4e00-\u9fa5][^\n\r]{5,}?){0,2}"  # 允许副标题
        r"(?![\.\d])",
        re.UNICODE
    )

    # 过滤规则增强（兼容OCR空格噪声）
    filter_keywords = re.compile(
        r"^(摘要|关键词|目录|作者|单位|日期|页码|图\d+|表\d+|.*比赛.*|.*泰迪杯.*)",
        re.IGNORECASE | re.UNICODE
    )

    # 候选标题评分体系
    candidates = defaultdict(int)
    for match in title_pattern.finditer(merged_text):
        candidate = match.group().strip()
        # 清洗OCR噪声
        candidate = re.sub(r"[\s\u3000]+", " ", candidate)  # 合并多余空白
        candidate = re.sub(r"^[^a-zA-Z\u4e00-\u9fa5]+", "", candidate)  # 去除开头非法字符

        if 12 <= len(candidate) <= 50 and not filter_keywords.match(candidate):
            # 关键词匹配评分
            keyword_score = sum(1 for kw in domain_keywords if kw in candidate)
            # 结构特征评分
            struct_score = 5 if re.search(r"[：\-—]", candidate) else 0
            # 综合评分
            candidates[candidate] += keyword_score * 3 + struct_score + len(candidate) / 10

    # 返回最高分标题
    if candidates:
        return max(candidates.items(), key=lambda x: x[1])[0]
    return "未识别到有效标题"


# --------------------------
# 集成调用接口
# --------------------------
def extract_title(pdf_path):
    """集成OCR的标题提取入口"""
    # Step 1: OCR识别PDF内容
    text = pdf_to_text(pdf_path)
    print("【扫描版PDF内容识别结果】")
    print(text)
    # Step 2: 从OCR文本提取标题
    return extract_title_from_ocr_text(text)


# --------------------------
# 测试用例
# --------------------------
if __name__ == "__main__":
    pdf_path = "附件3/B3108.pdf"  # 扫描版PDF路径
    title = extract_title(pdf_path)
    print("【扫描版PDF标题识别结果】")
    print(title)